import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from matplotlib.colors import LinearSegmentedColormap

# annualized_shape, annualized_return, annualized_volatility, turnover, margin, fitness, max_drawdown, \
#     information_ratio, calmar_ratio, trading_counts, winning_rate, profit_loss_ratio

for result_label in ["annualized_shape", "annualized_return", "annualized_volatility", "turnover", "margin",
                     "fitness", "max_drawdown", "information_ratio", "calmar_ratio", "trading_counts",
                     "winning_rate", "profit_loss_ratio"]:

    df = pd.read_csv('result/backtest_result_heatmap.csv')

    # 使用 pivot_table 转换数据
    pivot_df = df.pivot_table(
        values=f'{result_label}',
        index='oi_s',      # Y轴作为行
        columns='oi_l',    # X轴作为列
        aggfunc='mean'      # 处理重复值：取平均值
    )

    print("转换后的数据:")
    print(pivot_df)

    # 绘制热力图
    plt.figure(figsize=(10, 8))
    colors = ["green", "yellow", "red"]
    if result_label == "max_drawdown":
        colors = ["green", "yellow", "red"][::-1]
    cmap = LinearSegmentedColormap.from_list("rg", colors, N=256)
    sns.heatmap(
        pivot_df,
        annot=True,           # 显示数值
        fmt='.1f',           # 数值格式
        cmap=cmap,      # 颜色映射
        cbar_kws={'label': 'Value Scale'},
        linewidths=0.5,
        linecolor='white'
    )

    plt.title(f'{result_label}')
    plt.xlabel('M')
    plt.ylabel('N')
    plt.tight_layout()
    # plt.show()
    plt.savefig(f'./heatmap_png/{result_label}.png')